import sqlite3
from utils.debug import debug_wrap

class SqlitePKCache(object):
    """
    A cache for sqlite to map primary keys
    """
    def __init__(self, cursor, get_sql):
        self.cursor = cursor
        self.get_sql = get_sql
        self.cache = {}
        
    @debug_wrap
    def __getitem__(self, key):
        """
        Get the database key for a python key.
        """
        if key in self.cache:
            return self.cache[key]
        else:
            self.cursor.execute(self.get_sql, (key,))
            row = self.cursor.fetchone()
            if row:
                self.cache[key] = row[0]
                return row[0]
            else:
                raise KeyError, key
                
    @debug_wrap
    def __delitem__(self, key):
        """
        Delete an item. This is useful for example if you map game userids to
        database userids.
        """
        if key in self.cache:
            del self.cache[key]
            

class SimpleConnection(object):
    """
    A simple wrapper around sqlite3 connections and cursors.
    """
    def __init__(self, dbfile, pragmas=[('locking_mode', 'EXCLUSIVE'),
                                        ('synchronous', 'OFF'),
                                        ('journal_mode', 'OFF')]):
        """
        @param dbfile: Database file to use
        @param pragmas: PRAGMA statements to use after connecting to the
            database to speed stuff up.
        """
        self.connection = sqlite3.connect(dbfile)
        self.commit     = self.connection.commit
        self.cursor     = self.connection.cursor()
        self.execute    = self.cursor.execute
        self.fetchone   = self.cursor.fetchone
        self.fetchall   = self.cursor.fetchall
        for conf, value in pragmas:
            self.execute("PRAGMA %s=%s" % (conf, value))
        
    @debug_wrap
    def __call__(self, query, data=()):
        self.execute(query, data)
        
    @debug_wrap
    def smart_fetchone(self, sql, data=(), default=None):
        """
        Shortcut to execute a query and then get one row for this. It will
        return `default` if no rows are returned. If a row only contains one
        value, that value will be returned. Otherwise you'll get a tuple.
        @param sql: The SQL statement to execute.
        @param data: The data to insert into the statement.
        @param default: Default to return if no rows are found.
        """
        self.execute(sql, data)
        row = self.fetchone()
        if not row:
            return default
        if len(row) == 1:
            return row[0]
        return row
    
    @debug_wrap
    def smart_fetchall(self, sql, data=(), default=None):
        """
        Shortcut to execute a query and then get all rows for this. It will
        return `default` if no rows are returned. If a row only contains one
        value, that value will be returned. Otherwise you'll get a tuple.
        @param sql: The SQL statement to execute.
        @param data: The data to insert into the statement.
        @param default: Default to return if no rows are found.
        """
        self.execute(sql, data)
        rows = self.fetchall()
        if not rows:
            return default
        return [row if len(row) > 1 else row[0] for row in rows]
    
    def save(self):
        """
        Save the database.
        """
        self.commit()
        
    def close(self):
        """
        Close the connection to the database.
        """
        self.connection.close()